# -*- coding:utf-8 -*-
import hashlib
import base64
import hmac
import time
from urllib.parse import urlencode
import json
import websocket
import _thread as thread
import ssl

from ws.config.constants import xunfei

appid = xunfei.appid.value
api_secret = xunfei.api_secret.value

content = ""
class Document_Q_And_A:
    def __init__(self, APPId, APISecret, TimeStamp, OriginUrl):
        self.appId = APPId
        self.apiSecret = APISecret
        self.timeStamp = TimeStamp
        self.originUrl = OriginUrl

    def get_origin_signature(self):
        m2 = hashlib.md5()
        data = bytes(self.appId + self.timeStamp, encoding="utf-8")
        m2.update(data)
        checkSum = m2.hexdigest()
        return checkSum

    def get_signature(self):
        # 获取原始签名
        signature_origin = self.get_origin_signature()
        # print(signature_origin)
        # 使用加密键加密文本
        signature = hmac.new(self.apiSecret.encode('utf-8'), signature_origin.encode('utf-8'),
                             digestmod=hashlib.sha1).digest()
        # base64密文编码
        signature = base64.b64encode(signature).decode(encoding='utf-8')
        # print(signature)
        return signature

    def get_header(self):
        signature = self.get_signature()
        header = {
            "Content-Type": "application/json",
            "appId": self.appId,
            "timestamp": self.timeStamp,
            "signature": signature
        }
        return header

    def get_url(self):
        signature = self.get_signature()
        header = {
            "appId": self.appId,
            "timestamp": self.timeStamp,
            "signature": signature
        }
        return self.originUrl + "?" + f'appId={self.appId}&timestamp={self.timeStamp}&signature={signature}'
        # 使用urlencode会导致签名乱码
        # return self.originUrl + "?" + urlencode(header)

    def get_body(self, question):
        temp = historyMessage.copy()
        global lastUserMessage
        lastUserMessage = createMsg("user", question)
        temp.append(lastUserMessage)
        data = {
            "chatExtends": {
                "wikiPromptTpl": "请将以下内容作为已知信息：\n<wikicontent>\n请根据以上内容回答用户的问题。\n问题:<wikiquestion>\n回答:",
                # "wikiPromptTpl": "请将以下内容作为已知信息：\n<wikicontent>\n，将以上信息分段，每一段添加一个字母序号，从大写A开始进行输出，其它内容不变。",
                "wikiFilterScore": 72,
                "temperature": 0.5
            },
            "fileIds": [
                "05a3161ef2b447cfa4f2abb81350b95b",
                "a03fb96dc35848e8b51e4038ad0ba5fd",
                "2a3b27d4b7744b0686b0086c27e7c9e2",
                "087595775fbf47e6b2a0de44914322e5",
                "dfe250ea71b24770b875ae17d6196e44",
                "72b95428e6ab438cb37344e9489dbfed"
            ],
            "messages": temp
        }

        return data


# 收到websocket错误的处理
def on_error(ws, error):
    print("### error:", error)
    pass


# 收到websocket关闭的处理
def on_close(ws, close_status_code, close_msg):
    print("### closed ###")
    print("关闭代码：", close_status_code)
    print("关闭原因：", close_msg)
    pass

# 收到websocket连接建立的处理
def on_open(ws):
    thread.start_new_thread(run, (ws,))


def run(ws, *args):
    data = json.dumps(ws.question)
    ws.send(data)


# 收到websocket消息的处理
def on_message(ws, message):
    # print(message)
    data = json.loads(message)
    code = data['code']
    if code != 0:
        print(f'请求错误: {code}, {data}')
        websocketTemp.send(data['message'])
        ws.close()
    else:
        content = data["content"]
        status = data["status"]
        # print(f'status = {status}')
        # global answer
        # answer += "\t" + content
        print(content, end='')
        websocketTemp.send(content)
        if status == 2:
            lastAssistantMessage = createMsg("assistant", answer)
            historyMessage.append(lastUserMessage)
            historyMessage.append(lastAssistantMessage)
            ws.close()


def sendQuestion(webTest, question):
    global websocketTemp
    websocketTemp = webTest

    curTime = str(int(time.time()))
    OriginUrl = "wss://chatdoc.xfyun.cn/openapi/chat"
    document_Q_And_A = Document_Q_And_A(appid, api_secret, curTime, OriginUrl)

    wsUrl = document_Q_And_A.get_url()
    # print(wsUrl)
    headers = document_Q_And_A.get_header()
    body = document_Q_And_A.get_body(question)

    # 禁用WebSocket库的跟踪功能，使其不再输出详细的调试信息。
    websocket.enableTrace(False)
    ws = websocket.WebSocketApp(wsUrl, on_message=on_message, on_error=on_error, on_close=on_close, on_open=on_open)
    ws.appid = appid
    ws.question = body
    ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})


    # 文档问答成功

lastUserMessage = None
lastAssistantMessage = None
historyMessage = []
answer = ""
websocketTemp = ""
def createMsg(type, content):
    rst = {}
    rst["role"] = type
    rst["content"] = content
    return rst

# if __name__ == '__main__':
#     while (1):
#         answer = ""
#         Input = input("\n" + "我:")
#         print("小青:", end="")
#         sendQuestion(Input)

        # print(str(text))